# -*- coding: utf-8 -*-
import pandas as pd
from typing import Optional
import numpy as np
from collections import deque
from datetime import datetime
import time
from xtquant import xtdata
xtdata.enable_hello = False


class IndependentNineTurnCalculator:
    """
    独立日内九转计算器（参考VMA和RSI的模式）

    特点：
    1. 不依赖订阅，主动获取分钟级数据
    2. 使用缓存机制避免重复计算
    3. 每分钟更新一次，确保数据时效性
    4. 支持交易时段过滤
    """

    def __init__(self, compare_offset: int = 4, required_bars: int = 9,
                 max_window_minutes: int = 240, verbose: bool = True,
                 restrict_session: bool = True, xtdata=None):
        self.compare_offset = int(compare_offset)
        self.required_bars = int(required_bars)
        self.max_window_minutes = int(max_window_minutes)
        self.verbose = bool(verbose)
        self.restrict_session = bool(restrict_session)
        self.xtdata = xtdata

        # 交易时段定义
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()

        # 状态存储
        self._bar_state = {}  # code -> {minute:int, open:float, high:float, low:float, close:float}
        self._closes = {}  # code -> deque of close prices
        self._counts = {}  # code -> {"up":int, "down":int}
        self._last_signal = {}  # code -> last signal type
        self._last_signal_time = {}  # code -> last signal time

        # 缓存机制：避免重复下载同一天的数据
        self._cache_date = {}  # 股票代码 -> 缓存日期
        self._cache_ttl = 60  # 缓存有效期（秒）- 1分钟更新一次确保获取最新完整分钟数据
        self._cache_time = {}  # 股票代码 -> 缓存时间戳
        self._cache_minute = {}  # 股票代码 -> 缓存的分钟时间戳

    def reset(self, code: str = None):
        """重置指定股票或全部股票的状态"""
        if code is None:
            self._bar_state.clear()
            self._closes.clear()
            self._counts.clear()
            self._last_signal.clear()
            self._last_signal_time.clear()
            self._cache_date.clear()
            self._cache_time.clear()
            self._cache_minute.clear()
        else:
            self._bar_state.pop(code, None)
            self._closes.pop(code, None)
            self._counts.pop(code, None)
            self._last_signal.pop(code, None)
            self._last_signal_time.pop(code, None)
            self._cache_date.pop(code, None)
            self._cache_time.pop(code, None)
            self._cache_minute.pop(code, None)

    def _log(self, code: str, message: str):
        """日志记录"""
        if self.verbose:
            ##print(f"【独立九转】{code} {message}")
            pass

    def _is_trading_time(self, dt: datetime) -> bool:
        """检查是否在交易时段内"""
        if not self.restrict_session:
            return True

        tm = dt.time()
        # A股交易时间: 09:30-11:30, 13:00-15:00
        in_morning = (self.MORNING_START <= tm <= self.MORNING_END)
        in_afternoon = (self.AFTERNOON_START <= tm <= self.AFTERNOON_END)
        return in_morning or in_afternoon

    def _is_cache_valid(self, stock_code: str, today: str) -> bool:
        """检查缓存是否有效"""
        if stock_code not in self._cache_date or stock_code not in self._cache_time:
            return False

        # 检查日期是否匹配
        if self._cache_date[stock_code] != today:
            return False

        # 检查当前分钟是否与缓存分钟相同
        current_minute = datetime.now().strftime("%Y%m%d%H%M")
        cached_minute = self._cache_minute.get(stock_code, "")
        if cached_minute != current_minute:
            return False  # 分钟变化时重新计算

        # 检查时间是否在有效期内
        current_time = time.time()
        cache_time = self._cache_time[stock_code]
        return (current_time - cache_time) < self._cache_ttl

    def _get_minute_data(self, stock_code: str, today: str) -> pd.DataFrame:
        """获取分钟级数据"""
        try:
            # 检查xtdata是否可用
            if self.xtdata is None:
                self._log(stock_code, f"xtdata未初始化，无法获取数据")
                return pd.DataFrame()
            
            self._log(stock_code, f"开始下载 {today} 的分钟数据...")
            
            # 下载当日分钟数据
            self.xtdata.download_history_data(
                stock_code=stock_code,
                period="1m",
                start_time=today,
                end_time=today
            )
            time.sleep(1.5)  # 等待下载完成

            self._log(stock_code, f"开始获取本地数据...")
            
            # 获取本地数据
            data = self.xtdata.get_local_data(
                stock_list=[stock_code],
                period='1m',
                start_time=today,
                end_time=today
            )

            if not data:
                self._log(stock_code, f"get_local_data返回空值")
                return pd.DataFrame()
                
            if stock_code not in data:
                self._log(stock_code, f"股票代码 {stock_code} 不在返回数据中")
                self._log(stock_code, f"可用的股票代码: {list(data.keys())}")
                return pd.DataFrame()
                
            if data[stock_code] is None:
                self._log(stock_code, f"股票 {stock_code} 的数据为None")
                return pd.DataFrame()
                
            if data[stock_code].empty:
                self._log(stock_code, f"股票 {stock_code} 的数据为空DataFrame")
                return pd.DataFrame()

            df = data[stock_code].copy()
            self._log(stock_code, f"成功获取数据，原始列: {list(df.columns)}")
            self._log(stock_code, f"数据形状: {df.shape}")

            # 识别时间列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            if not time_col:
                self._log(stock_code, f"无法识别时间列，可用列: {list(df.columns)}")
                return pd.DataFrame()

            # 统一时间列格式
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')

            df = df[df['__time__'].notna()].sort_values('__time__')
            self._log(stock_code, f"时间列处理完成，有效数据: {len(df)} 条")

            # 识别价格列（优先使用close，其次使用其他价格列）
            price_cols = ['close', 'lastPrice', 'price', 'last', 'current']
            price_col = None
            for col in price_cols:
                if col in df.columns:
                    price_col = col
                    break

            if not price_col:
                self._log(stock_code, f"无法识别价格列，可用列: {list(df.columns)}")
                return pd.DataFrame()

            # 重命名价格列为统一格式
            df['price'] = df[price_col]
            
            result_df = df[['__time__', 'price']]
            self._log(stock_code, f"数据处理完成，最终数据: {len(result_df)} 条")
            
            return result_df

        except Exception as e:
            self._log(stock_code, f"获取分钟数据异常: {e}")
            import traceback
            self._log(stock_code, f"异常详情: {traceback.format_exc()}")
            return pd.DataFrame()

    def _filter_trading_session_data(self, df: pd.DataFrame, cutoff_time: datetime = None) -> pd.DataFrame:
        """筛选交易时段内的数据，并排除指定时间后未完成的分钟"""
        try:
            # 先筛选交易时段内的数据
            mask = [self._is_trading_time(ts) for ts in df['__time__']]
            if any(mask):
                filtered_df = df[mask]

                # 如果没有指定截止时间，使用当前时间
                if cutoff_time is None:
                    cutoff_time = datetime.now()
                
                # 排除未完成的分钟：只保留完整的分钟数据
                cutoff_minute = cutoff_time.replace(second=0, microsecond=0)

                # 只保留时间小于截止分钟的数据（即已完成的分钟）
                complete_mask = filtered_df['__time__'] < cutoff_minute
                if any(complete_mask):
                    return filtered_df[complete_mask]
        except Exception:
            pass
        return df

    def _process_minute_data(self, stock_code: str, df: pd.DataFrame, end_time: datetime = None) -> str:
        """处理分钟数据，计算九转信号"""
        if df.empty:
            return None

        # 初始化状态
        cnt = self._counts.setdefault(stock_code, {"up": 0, "down": 0})
        closes = self._closes.setdefault(stock_code, deque(maxlen=1000))

        signal = None
        last_signal = None  # 记录最后一个信号
        previous_minute_signal = None  # 记录上一分钟的信号

        # 按时间顺序处理每个分钟K线
        for _, row in df.iterrows():
            # 如果指定了结束时间，只处理该时间之前的数据
            if end_time is not None and row['__time__'] > end_time:
                break
                
            minute_bucket = int(row['__time__'].timestamp() // 60)
            price = row['price']

            if pd.isna(price):
                continue

            bar = self._bar_state.get(stock_code)

            # 新分钟K线开始
            if bar is None or minute_bucket != bar["minute"]:
                if bar is not None:
                    # 完成上一分钟K线
                    closes.append(bar["close"])

                    # 检查信号
                    signal = self._check_signal(stock_code, closes, cnt, row['__time__'])
                    if signal:
                        last_signal = signal
                        # 记录上一分钟的信号
                        previous_minute_signal = signal

                # 开始新K线
                self._bar_state[stock_code] = {
                    "minute": minute_bucket,
                    "open": float(price),
                    "high": float(price),
                    "low": float(price),
                    "close": float(price)
                }
            else:
                # 更新当前K线
                bar["close"] = float(price)
                if price > bar["high"]:
                    bar["high"] = float(price)
                if price < bar["low"]:
                    bar["low"] = float(price)

        # 重要：处理最后一个K线，确保信号不丢失
        if self._bar_state.get(stock_code) is not None:
            last_bar = self._bar_state[stock_code]
            closes.append(last_bar["close"])
            signal = self._check_signal(stock_code, closes, cnt, datetime.now()) # 使用当前时间
            if signal:
                last_signal = signal
                # 更新上一分钟的信号
                previous_minute_signal = signal

        # 返回上一分钟的信号，如果没有则返回None
        return previous_minute_signal

    def _check_signal(self, code: str, closes: deque, cnt: dict, current_time: datetime) -> str:
        """检查是否满足信号条件"""
        if len(closes) <= self.compare_offset:
            return None

        cur = closes[-1]
        ref = closes[-1 - self.compare_offset]

        # 更新计数：与同花顺更一致的容忍度
        # - 满足 cur >= ref 则看涨计数 +1，否则清零
        # - 满足 cur <= ref 则看跌计数 +1，否则清零
        # 注意：相等会同时累加两个方向，允许重叠序列
        if cur >= ref:
            cnt["up"] += 1
        else:
            cnt["up"] = 0

        if cur <= ref:
            cnt["down"] += 1
        else:
            cnt["down"] = 0

        # 检查信号
        signal = None
        if cnt["up"] >= self.required_bars:
            signal = "卖出"
            # 仅重置对应方向，允许另一方向的序列延续
            cnt["up"] = 0
            self._log(code, f"触发卖出信号，连续{self.required_bars}次收盘价不低于前{self.compare_offset}根K线")

        elif cnt["down"] >= self.required_bars:
            signal = "买入"
            cnt["down"] = 0
            self._log(code, f"触发买入信号，连续{self.required_bars}次收盘价不高于前{self.compare_offset}根K线")

        if signal:
            self._last_signal[code] = signal
            # 记录信号触发时间（使用当前处理的时间，而不是系统时间）
            self._last_signal_time[code] = current_time  # 这里需要传入实际的时间
            self._log(code, f"信号触发: {signal}, 当前价格: {cur:.4f}, 参考价格: {ref:.4f}")

        return signal

    def get_latest_signal(self, stock_code: str) -> str:
        """获取最新的九转信号"""
        try:
            today = datetime.now().strftime("%Y%m%d")

            # 检查缓存是否有效
            if self._is_cache_valid(stock_code, today):
                return self._last_signal.get(stock_code, "")

            # 获取分钟数据
            df = self._get_minute_data(stock_code, today)
            if df.empty:
                return ""

            # 筛选交易时段数据
            filtered_df = self._filter_trading_session_data(df)
            if filtered_df.empty:
                return ""

            # 处理数据并计算信号
            signal = self._process_minute_data(stock_code, filtered_df, None)

            # 更新缓存
            if signal is not None:
                self._cache_date[stock_code] = today
                self._cache_time[stock_code] = time.time()
                self._cache_minute[stock_code] = datetime.now().strftime("%Y%m%d%H%M")
                self._log(stock_code, f"信号计算完成: {signal}")
            else:
                # 即使没有信号也要更新缓存
                self._cache_date[stock_code] = today
                self._cache_time[stock_code] = time.time()
                self._cache_minute[stock_code] = datetime.now().strftime("%Y%m%d%H%M")

            return signal or ""

        except Exception as e:
            self._log(stock_code, f"获取最新信号异常: {e}")
            return ""

    def _filter_data_before_time(self, df: pd.DataFrame, end_time: datetime) -> pd.DataFrame:
        """
        过滤掉指定时间之后的数据（回测模式使用）
        
        Args:
            df: 原始数据DataFrame
            end_time: 结束时间
            
        Returns:
            过滤后的DataFrame
        """
        try:
            if df.empty:
                return df
            
            # 确保时间列存在
            if '__time__' not in df.columns:
                return df
            
            # 过滤掉结束时间之后的数据
            filtered_df = df[df['__time__'] <= end_time]
            return filtered_df
            
        except Exception as e:
            self._log("", f"数据过滤异常: {e}")
            return df

    def get_previous_minute_signal(self, stock_code: str, end_time: Optional[datetime] = None) -> str:
        """获取上一分钟的完整九转信号（推荐使用）"""
        try:
            if end_time is not None:
                # 回测模式：使用指定的结束时间
                today = end_time.strftime("%Y%m%d")
                current_time = end_time
                self._log(stock_code, f"回测模式，目标时间: {end_time}")
            else:
                # 实时模式：使用当前时间
                today = datetime.now().strftime("%Y%m%d")
                current_time = datetime.now()
                self._log(stock_code, f"实时模式，当前时间: {current_time}")
            
            # 计算上一分钟的时间
            previous_minute = current_time.replace(second=0, microsecond=0) - pd.Timedelta(minutes=1)
            previous_minute_str = previous_minute.strftime("%Y%m%d%H%M")
            self._log(stock_code, f"上一分钟时间: {previous_minute}")
            
            # 检查是否有上一分钟的缓存信号（回测模式下不使用缓存）
            if end_time is None and (stock_code in self._cache_minute and 
                self._cache_minute[stock_code] == previous_minute_str and
                stock_code in self._last_signal):
                return self._last_signal[stock_code]
            
            # 获取分钟数据
            df = self._get_minute_data(stock_code, today)
            if df.empty:
                return ""

            self._log(stock_code, f"原始数据时间范围: {df['__time__'].min()} 到 {df['__time__'].max()}")

            # 在回测模式下，直接使用end_time作为截止时间
            if end_time is not None:
                # 回测模式：先过滤掉结束时间之后的数据
                filtered_df = self._filter_data_before_time(df, end_time)
                if filtered_df.empty:
                    self._log(stock_code, f"过滤后数据为空")
                    return ""
                
                self._log(stock_code, f"过滤后数据时间范围: {filtered_df['__time__'].min()} 到 {filtered_df['__time__'].max()}")
                
                # 再筛选交易时段数据
                filtered_df = self._filter_trading_session_data(filtered_df, cutoff_time=end_time)
                if filtered_df.empty:
                    self._log(stock_code, f"交易时段过滤后数据为空")
                    return ""
                    
                self._log(stock_code, f"最终数据时间范围: {filtered_df['__time__'].min()} 到 {filtered_df['__time__'].max()}")
            else:
                # 实时模式：筛选交易时段数据，但只保留到上一分钟
                filtered_df = self._filter_trading_session_data(df, cutoff_time=previous_minute)
                if filtered_df.empty:
                    return ""
            
            # 重要：使用与测试代码相同的逻辑来计算信号，避免状态重置问题
            # 重置状态以重新计算
            self.reset(stock_code)
            self._log(stock_code, f"重置状态以确保计算准确性")
            
            # 使用与测试代码相同的逻辑来计算信号
            signal_times = []
            cnt = self._counts.setdefault(stock_code, {"up": 0, "down": 0})
            closes = self._closes.setdefault(stock_code, deque(maxlen=1000))
            
            # 按时间排序
            filtered_df = filtered_df.sort_values('__time__')
            
            # 使用与测试代码相同的循环逻辑
            for _, row in filtered_df.iterrows():
                minute_bucket = int(row['__time__'].timestamp() // 60)
                price = row['price']
                
                if pd.isna(price):
                    continue
                
                bar = self._bar_state.get(stock_code)
                
                # 新分钟K线开始
                if bar is None or minute_bucket != bar["minute"]:
                    if bar is not None:
                        # 完成上一分钟K线
                        closes.append(bar["close"])
                        
                        # 检查信号
                        if len(closes) > self.compare_offset:
                            cur = closes[-1]
                            ref = closes[-1 - self.compare_offset]
                            
                            # 更新计数
                            if cur >= ref:
                                cnt["up"] += 1
                                cnt["down"] = 0
                            else:
                                cnt["up"] = 0
                                cnt["down"] += 1
                            
                            # 检查信号
                            if cnt["up"] >= self.required_bars:
                                signal_times.append({
                                    'time': row['__time__'],
                                    'signal': '卖出',
                                    'price': price
                                })
                                cnt["up"] = 0
                                
                            elif cnt["down"] >= self.required_bars:
                                signal_times.append({
                                    'time': row['__time__'],
                                    'signal': '买入',
                                    'price': price
                                })
                                cnt["down"] = 0
                    
                    # 开始新K线
                    self._bar_state[stock_code] = {
                        "minute": minute_bucket,
                        "open": float(price),
                        "high": float(price),
                        "low": float(price),
                        "close": float(price)
                    }
                else:
                    # 更新当前K线
                    bar["close"] = float(price)
                    if price > bar["high"]:
                        bar["high"] = float(price)
                    if price < bar["low"]:
                        bar["low"] = float(price)
            
            # 找到上一分钟对应的信号
            signal = ""
            if signal_times:
                # 找到在上一分钟时间范围内的最后一个信号
                for signal_info in reversed(signal_times):
                    signal_time = signal_info['time']
                    if previous_minute <= signal_time < current_time:
                        signal = signal_info['signal']
                        break
            
            # 更新缓存为上一分钟（回测模式下不更新缓存）
            if end_time is None:
                self._cache_date[stock_code] = today
                self._cache_time[stock_code] = time.time()
                self._cache_minute[stock_code] = previous_minute_str
            
            if signal:
                self._log(stock_code, f"上一分钟信号: {signal}")
            
            return signal or ""

        except Exception as e:
            self._log(stock_code, f"获取上一分钟信号异常: {e}")
            return ""

    def get_nine_turn_analysis(self, stock_code: str) -> dict:
        """获取九转指标分析"""
        try:
            today = datetime.now().strftime("%Y%m%d")

            # 检查缓存是否有效
            if self._is_cache_valid(stock_code, today):
                return {
                    'last_signal': self._last_signal.get(stock_code, ""),
                    'up_count': self._counts.get(stock_code, {}).get('up', 0),
                    'down_count': self._counts.get(stock_code, {}).get('down', 0),
                    'closes_count': len(self._closes.get(stock_code, [])),
                    'cached': True
                }

            # 获取最新信号（这会更新缓存）
            signal = self.get_latest_signal(stock_code)

            return {
                'last_signal': signal,
                'up_count': self._counts.get(stock_code, {}).get('up', 0),
                'down_count': self._counts.get(stock_code, {}).get('down', 0),
                'closes_count': len(self._closes.get(stock_code, [])),
                'cached': False
            }

        except Exception as e:
            self._log(stock_code, f"获取九转分析异常: {e}")
            return {
                'last_signal': "",
                'up_count': 0,
                'down_count': 0,
                'closes_count': 0,
                'cached': False,
                'error': str(e)
            }

    def check_nine_turn_conditions(self, stock_code: str) -> dict:
        """检查九转条件（兼容原有接口）"""
        signal = self.get_latest_signal(stock_code)
        analysis = self.get_nine_turn_analysis(stock_code)

        return {
            'is_buy': signal == "买入",
            'is_sell': signal == "卖出",
            'signal': signal,
            'analysis': analysis,
            'message': f"九转信号: {signal}, 上涨计数: {analysis.get('up_count', 0)}, 下跌计数: {analysis.get('down_count', 0)}"
        }

    # 向后兼容的方法
    def on_tick(self, code: str, timestamp, price) -> str:
        """兼容原有接口，但现在建议使用get_latest_signal"""
        return self.get_latest_signal(code)


# 保持向后兼容的别名
THSNineTurn = IndependentNineTurnCalculator

# 示例用法和测试
if __name__ == "__main__":
    print("开始测试独立九转计算器...")

    # 初始化独立九转计算器
    nine_turn = IndependentNineTurnCalculator(
        compare_offset=4,
        required_bars=9,
        verbose=True,
        restrict_session=True,
        xtdata=xtdata
    )

    # 股票代码
    code = "159268.SZ"
    print(f"测试股票: {code}")

    # 测试获取上一分钟信号

    try:
        # 使用固定的测试时间
        print("\n🔍 使用固定的测试时间")
        target_time = datetime(2025, 8, 22, 9, 57, 1)
        
        print(f"目标时间: {target_time}")
        
        signal = nine_turn.get_previous_minute_signal(code, target_time)
        print(f"✅ 上一分钟九转信号: {signal}")
        
    except Exception as e:
        print(f"❌ 获取上一分钟信号失败: {e}")

    # 分析所有买入卖出信号时间
    print("\n🔍 分析所有买入卖出信号时间")
    try:
        ##today = datetime.now().strftime("%Y%m%d")
        today = "20250822"
        
        # 获取分钟数据
        df = nine_turn._get_minute_data(code, today)
        
        if not df.empty:
            # 筛选交易时段数据
            filtered_df = nine_turn._filter_trading_session_data(df)
            
            if not filtered_df.empty:
                # 使用与回测模式相同的逻辑来计算信号
                signal_times = []
                
                # 重置状态以重新计算
                nine_turn.reset(code)
                
                # 按时间排序
                filtered_df = filtered_df.sort_values('__time__')
                
                # 使用九转计算器的核心方法来计算信号
                for _, row in filtered_df.iterrows():
                    minute_bucket = int(row['__time__'].timestamp() // 60)
                    price = row['price']
                    
                    if pd.isna(price):
                        continue
                    
                    bar = nine_turn._bar_state.get(code)
                    
                    # 新分钟K线开始
                    if bar is None or minute_bucket != bar["minute"]:
                        if bar is not None:
                            # 完成上一分钟K线
                            closes = nine_turn._closes.setdefault(code, deque(maxlen=1000))
                            closes.append(bar["close"])
                            
                            # 检查信号
                            if len(closes) > nine_turn.compare_offset:
                                cur = closes[-1]
                                ref = closes[-1 - nine_turn.compare_offset]
                                
                                cnt = nine_turn._counts.setdefault(code, {"up": 0, "down": 0})
                                
                                # 更新计数
                                if cur >= ref:
                                    cnt["up"] += 1
                                    cnt["down"] = 0
                                else:
                                    cnt["up"] = 0
                                    cnt["down"] += 1
                                
                                # 检查信号
                                if cnt["up"] >= nine_turn.required_bars:
                                    signal_times.append({
                                        'time': row['__time__'],
                                        'signal': '卖出',
                                        'price': price
                                    })
                                    cnt["up"] = 0
                                    
                                elif cnt["down"] >= nine_turn.required_bars:
                                    signal_times.append({
                                        'time': row['__time__'],
                                        'signal': '买入',
                                        'price': price
                                    })
                                    cnt["down"] = 0
                        
                        # 开始新K线
                        nine_turn._bar_state[code] = {
                            "minute": minute_bucket,
                            "open": float(price),
                            "high": float(price),
                            "low": float(price),
                            "close": float(price)
                        }
                    else:
                        # 更新当前K线
                        bar["close"] = float(price)
                        if price > bar["high"]:
                            bar["high"] = float(price)
                        if price < bar["low"]:
                            bar["low"] = float(price)
                
                # 输出所有信号时间
                if signal_times:
                    print(f"\n⏰ 九转信号时间点位:")
                    for i, signal_info in enumerate(signal_times, 1):
                        time_str = signal_info['time'].strftime('%H:%M:%S')
                        signal_type = signal_info['signal']
                        price = signal_info['price']
                        print(f"   {i:2d}. {time_str} | {signal_type:4s} | 价格: {price:7.4f}")
                else:
                    print("   当前数据中未发现九转信号")
                    
        else:
            print("   未获取到分钟数据")
            
    except Exception as e:
        print(f"❌ 分析信号时间失败: {e}")
        import traceback
        traceback.print_exc()
    
